module state_mod

	use const_mod
	use namelist_mod
	use mesh_mod
	use allocator_mod

	implicit none

	private
	public state_type
	public state

	type state_type
		type(mesh_type), pointer :: mesh => null()
		real(r8), allocatable, dimension(:,:,:) :: u
		real(r8), allocatable, dimension(:,:,:) :: v
		real(r8), allocatable, dimension(:,:,:) :: we     ! ∂π/∂η * \dot(η)
		real(r8), allocatable, dimension(:,:  ) :: phs    ! surface hpressure
		real(r8), allocatable, dimension(:,:,:) :: ph_lev ! hpressure at interface
		real(r8), allocatable, dimension(:,:,:) :: ph     ! hpressure at layer
		real(r8), allocatable, dimension(:,:,:) :: m      ! layer mass
		real(r8), allocatable, dimension(:,:,:) :: m_lon  ! mass on zonal edge
		real(r8), allocatable, dimension(:,:,:) :: m_lat  ! mass on meridional edge
		real(r8), allocatable, dimension(:,:,:) :: m_lev  ! mass on interface edge
		real(r8), allocatable, dimension(:,:,:) :: mf_lon ! mass flux on zonal edge
		real(r8), allocatable, dimension(:,:,:) :: mf_lat ! mass flux on meridional edge
		real(r8), allocatable, dimension(:,:  ) :: psi    ! streamfunction for nondivergence wind in 2D adv test
		real(r8), allocatable, dimension(:,:  ) :: zs     ! surface elevation
  contains
  	procedure :: init  => state_init
  	procedure :: clear => state_clear
  	final :: state_final
	end type state_type

  type(state_type), allocatable :: state(:)
contains
	
	subroutine state_init(this, mesh)

		class(state_type), intent(inout) :: this
		type(mesh_type), intent(in) :: mesh

		call this%clear()

		call allocate_array(mesh, this%u     , half_lon=.true., full_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%v     , full_lon=.true., half_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%we    , full_lon=.true., full_lat=.true., half_lev=.true.)
		call allocate_array(mesh, this%phs   , full_lon=.true., full_lat=.true.                 )
		call allocate_array(mesh, this%ph_lev, full_lon=.true., full_lat=.true., half_lev=.true.)
		call allocate_array(mesh, this%ph    , full_lon=.true., full_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%m     , full_lon=.true., full_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%m_lon , half_lon=.true., full_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%m_lat , full_lon=.true., half_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%m_lev , full_lon=.true., full_lat=.true., half_lev=.true.)
		call allocate_array(mesh, this%mf_lon, half_lon=.true., full_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%mf_lat, full_lon=.true., half_lat=.true., full_lev=.true.)
		call allocate_array(mesh, this%psi   , half_lon=.true., half_lat=.true.								  )
		call allocate_array(mesh, this%zs    , full_lon=.true., full_lat=.true.								  )

	end subroutine state_init

	subroutine state_clear(this)

		class(state_type), intent(inout) :: this

		if (allocated(this%u      )) deallocate(this%u     )
		if (allocated(this%v      )) deallocate(this%v     )
		if (allocated(this%we     )) deallocate(this%we    )
		if (allocated(this%ph_lev )) deallocate(this%ph_lev)
		if (allocated(this%ph     )) deallocate(this%ph    )
		if (allocated(this%m      )) deallocate(this%m     )
		if (allocated(this%m_lon  )) deallocate(this%m_lon )
		if (allocated(this%m_lat  )) deallocate(this%m_lat )
		if (allocated(this%m_lev  )) deallocate(this%m_lev )
		if (allocated(this%mf_lon )) deallocate(this%mf_lon)
		if (allocated(this%mf_lat )) deallocate(this%mf_lat)
		if (allocated(this%psi    )) deallocate(this%psi   )
		if (allocated(this%zs     )) deallocate(this%zs    )

	end subroutine state_clear

	subroutine state_final(this)
	  
	  type(state_type), intent(inout) :: this

	  call this%clear()
	    
	end subroutine state_final

end module state_mod